package com.ctfo.ulcp.exp.core;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import io.weaviate.client.Config;
import io.weaviate.client.WeaviateAuthClient;
import io.weaviate.client.WeaviateClient;
import io.weaviate.client.base.Result;
import io.weaviate.client.base.WeaviateErrorMessage;
import io.weaviate.client.v1.auth.exception.AuthException;
import io.weaviate.client.v1.batch.model.BatchDeleteResponse;
import io.weaviate.client.v1.batch.model.ObjectGetResponse;
import io.weaviate.client.v1.data.model.WeaviateObject;
import io.weaviate.client.v1.filters.WhereFilter;
import io.weaviate.client.v1.graphql.model.GraphQLError;
import io.weaviate.client.v1.graphql.model.GraphQLResponse;
import io.weaviate.client.v1.graphql.query.Get;
import io.weaviate.client.v1.graphql.query.argument.NearVectorArgument;
import io.weaviate.client.v1.graphql.query.argument.WhereArgument;
import io.weaviate.client.v1.graphql.query.fields.Field;
import lombok.Builder;

import java.nio.charset.StandardCharsets;
import java.util.*;

import static dev.langchain4j.internal.Utils.*;
import static dev.langchain4j.internal.ValidationUtils.ensureNotBlank;
import static io.weaviate.client.v1.data.replication.model.ConsistencyLevel.QUORUM;
import static java.util.Arrays.stream;
import static java.util.Collections.emptyList;
import static java.util.Collections.singletonList;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

/**
 * Represents the <a href="https://weaviate.io/">Weaviate</a> vector database.
 * Current implementation assumes the cosine distance metric is used.
 * Does not support storing {@link dev.langchain4j.data.document.Metadata} yet.
 */
public class WeaviateEmbeddingStore implements CrudEmbeddingStore {

    private static final String METADATA_TEXT_SEGMENT = "text";
    private static final String ADDITIONALS = "_additional";

    private final WeaviateClient client;
    private final String objectClass;
    private final boolean avoidDups;
    private final String consistencyLevel;

    /**
     * 指明哪些元数据属性在查询时候要返回
     */
    private final List<String> metadataKeys;

    /**
     * Creates a new WeaviateEmbeddingStore instance.
     *
     * @param apiKey           Your Weaviate API key. Not required for local deployment.
     * @param scheme           The scheme, e.g. "https" of cluster URL. Find in under Details of your Weaviate cluster.
     * @param host             The host, e.g. "langchain4j-4jw7ufd9.weaviate.network" of cluster URL.
     *                         Find in under Details of your Weaviate cluster.
     * @param port             The port, e.g. 8080. This parameter is optional.
     * @param objectClass      The object class you want to store, e.g. "MyGreatClass". Must start from an uppercase letter.
     * @param avoidDups        If true (default), then <code>WeaviateEmbeddingStore</code> will generate a hashed ID based on
     *                         provided text segment, which avoids duplicated entries in DB.
     *                         If false, then random ID will be generated.
     * @param consistencyLevel Consistency level: ONE, QUORUM (default) or ALL. Find more details <a href="https://weaviate.io/developers/weaviate/concepts/replication-architecture/consistency#tunable-write-consistency">here</a>.
     */
    @Builder
    public WeaviateEmbeddingStore(
            String apiKey,
            String scheme,
            String host,
            Integer port,
            String objectClass,
            Boolean avoidDups,
            String consistencyLevel,
            List<String> metadataKeys
    ) {
        try {

            Config config = new Config(
                    ensureNotBlank(scheme, "scheme"),
                    concatenate(ensureNotBlank(host, "host"), port)
            );
            this.client = WeaviateAuthClient.apiKey(config, getOrDefault(apiKey, ""));
        } catch (AuthException e) {
            throw new IllegalArgumentException(e);
        }
        this.objectClass = getOrDefault(objectClass, "Default");
        this.avoidDups = getOrDefault(avoidDups, true);
        this.consistencyLevel = getOrDefault(consistencyLevel, QUORUM);
        this.metadataKeys = getOrDefault(metadataKeys, Collections.emptyList());
    }

    private static String concatenate(String host, Integer port) {
        if (port == null) {
            return host;
        } else {
            return host + ":" + port;
        }
    }


    @Override
    public List<String> specialMetaKeys() {
        return metadataKeys;
    }

    @Override
    public String add(Embedding embedding) {
        String id = randomUUID();
        add(id, embedding);
        return id;
    }

    /**
     * Adds a new embedding with provided ID to the store.
     *
     * @param id        the ID of the embedding to add in UUID format, since it's Weaviate requirement.
     *                  See <a href="https://weaviate.io/developers/weaviate/manage-data/create#id">Weaviate docs</a> and
     *                  <a href="https://en.wikipedia.org/wiki/Universally_unique_identifier">UUID on Wikipedia</a>
     * @param embedding the embedding to add
     */
    @Override
    public void add(String id, Embedding embedding) {
        addAll(singletonList(id), singletonList(embedding), null);
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        return addAll(singletonList(embedding), singletonList(textSegment)).stream().findFirst().orElse(null);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        return addAll(embeddings, null);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> embedded) {
        return addAll(null, embeddings, embedded);
    }

    // 扩展: 根据元数据字段查询
    public List<EmbeddingMatch<TextSegment>> findRelevant(
            Embedding referenceEmbedding,
            Map<String, String> meta,
            int maxResults,
            double minCertainty
    ) {

        WhereArgument whereArgument = null;
        if (meta != null && !meta.isEmpty()) {
            WhereFilter[] filters = meta.entrySet().stream().map(e -> WhereFilter.builder().path(e.getKey()).operator("Equal").valueText(e.getValue()).build()).toArray(WhereFilter[]::new);

            whereArgument = WhereArgument.builder().filter(
                    WhereFilter.builder().operator("And").operands(filters).build()
            ).build();
        }

        Get get = client
                .graphQL()
                .get()
                .withClassName(objectClass)
                .withFields(getQueryFields().toArray(new Field[0]))
                .withNearVector(
                        NearVectorArgument
                                .builder()
                                .vector(referenceEmbedding.vectorAsList().toArray(new Float[0]))
                                .certainty((float) minCertainty)
                                .build()
                )
                .withLimit(maxResults);

        // 增加查询条件
        if (whereArgument != null) {
            get.withWhere(whereArgument);
        }

        Result<GraphQLResponse> result = get.run();
        if (result.hasErrors()) {
            throw new IllegalArgumentException(
                    result.getError().getMessages().stream().map(WeaviateErrorMessage::getMessage).collect(joining("\n"))
            );
        }

        GraphQLError[] errors = result.getResult().getErrors();
        if (errors != null && errors.length > 0) {
            throw new IllegalArgumentException(stream(errors).map(GraphQLError::getMessage).collect(joining("\n")));
        }

        Optional<Map.Entry<String, Map>> resGetPart =
                ((Map<String, Map>) result.getResult().getData()).entrySet().stream().findFirst();
        if (!resGetPart.isPresent()) {
            return emptyList();
        }

        Optional resItemsPart = resGetPart.get().getValue().entrySet().stream().findFirst();
        if (!resItemsPart.isPresent()) {
            return emptyList();
        }

        List<Map<String, ?>> resItems = ((Map.Entry<String, List<Map<String, ?>>>) resItemsPart.get()).getValue();

        return resItems.stream().map(WeaviateEmbeddingStore::toEmbeddingMatch).collect(toList());
    }


    /**
     * 删除数据
     *
     * @param meta 为空的场合删除所有
     */
    public void batchDelete(Map<String, String> meta) {
        if (meta != null && !meta.isEmpty()) {
            WhereFilter[] filters = meta.entrySet().stream().map(e -> WhereFilter.builder().path(e.getKey()).operator("Equal").valueText(e.getValue()).build()).toArray(WhereFilter[]::new);
            WhereFilter whereFilter = WhereFilter.builder().operator("And").operands(filters).build();

            Result<BatchDeleteResponse> result = client.batch().objectsBatchDeleter()
                    .withClassName(objectClass)
                    .withWhere(whereFilter)
                    .run();
            if (result.hasErrors()) {
                throw new IllegalArgumentException(
                        result.getError().getMessages().stream().map(WeaviateErrorMessage::getMessage).collect(joining("\n"))
                );
            }
        } else {
            Result<Boolean> result = client.schema().classDeleter().withClassName(objectClass).run();
            if (result.hasErrors()) {
                throw new IllegalArgumentException(
                        result.getError().getMessages().stream().map(WeaviateErrorMessage::getMessage).collect(joining("\n"))
                );
            }
        }
    }

    /**
     * {@inheritDoc}
     * The score inside {@link EmbeddingMatch} is Weaviate's certainty.
     */
    @Override
    public List<EmbeddingMatch<TextSegment>> findRelevant(
            Embedding referenceEmbedding,
            int maxResults,
            double minCertainty
    ) {

        Result<GraphQLResponse> result = client
                .graphQL()
                .get()
                .withClassName(objectClass)
                .withFields(getQueryFields().toArray(new Field[0]))
                .withNearVector(
                        NearVectorArgument
                                .builder()
                                .vector(referenceEmbedding.vectorAsList().toArray(new Float[0]))
                                .certainty((float) minCertainty)
                                .build()
                )
                .withLimit(maxResults)
                .run();

        if (result.hasErrors()) {
            throw new IllegalArgumentException(
                    result.getError().getMessages().stream().map(WeaviateErrorMessage::getMessage).collect(joining("\n"))
            );
        }

        GraphQLError[] errors = result.getResult().getErrors();
        if (errors != null && errors.length > 0) {
            throw new IllegalArgumentException(stream(errors).map(GraphQLError::getMessage).collect(joining("\n")));
        }

        Optional<Map.Entry<String, Map>> resGetPart =
                ((Map<String, Map>) result.getResult().getData()).entrySet().stream().findFirst();
        if (!resGetPart.isPresent()) {
            return emptyList();
        }

        Optional resItemsPart = resGetPart.get().getValue().entrySet().stream().findFirst();
        if (!resItemsPart.isPresent()) {
            return emptyList();
        }

        List<Map<String, ?>> resItems = ((Map.Entry<String, List<Map<String, ?>>>) resItemsPart.get()).getValue();

        return resItems.stream().map(WeaviateEmbeddingStore::toEmbeddingMatch).collect(toList());
    }

    private List<Field> getQueryFields() {
        List<Field> fields = new ArrayList<>();
        fields.add(Field.builder().name(METADATA_TEXT_SEGMENT).build());
        fields.add(Field.builder().name(ADDITIONALS).fields(
                Field.builder().name("id").build(),
                Field.builder().name("certainty").build(),
                Field.builder().name("vector").build()
        ).build());
        for (String metaKey : specialMetaKeys()) {
            fields.add(Field.builder().name(metaKey).build());
        }
        return fields;
    }

    private List<String> addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> embedded) {
        if (embedded != null && embeddings.size() != embedded.size()) {
            throw new IllegalArgumentException("The list of embeddings and embedded must have the same size");
        }

        List<String> resIds = new ArrayList<>();
        List<WeaviateObject> objects = new ArrayList<>();
        for (int i = 0; i < embeddings.size(); i++) {
            // modify by cgy: 存储元数据
            String text = "";
            Metadata metadata = null;
            if (embedded != null) {
                TextSegment segment = embedded.get(i);
                text = segment.text();
                metadata = segment.metadata();
            }


            String id;
            if (ids != null) {
                id = ids.get(i);
            } else if (avoidDups && embedded != null) {
                id = generateUUIDFrom(embedded.get(i).text());
            } else if (metadata != null && metadataKeys != null && !metadataKeys.isEmpty()) {
                StringBuilder builder = new StringBuilder();
                for (String key : metadataKeys) {
                    String value = metadata.get(key);
                    builder.append(value != null ? value : "").append(":");
                }
                id = UUID.nameUUIDFromBytes(builder.toString().getBytes(StandardCharsets.UTF_8)).toString();
            } else {
                id = randomUUID();
            }
            resIds.add(id);

            objects.add(buildObject(id, embeddings.get(i), text, metadata));

        }

        Result<ObjectGetResponse[]> result = client
                .batch()
                .objectsBatcher()
                .withObjects(objects.toArray(new WeaviateObject[0]))
                .withConsistencyLevel(consistencyLevel)
                .run();

        // modify by cgy: 添加数据失败的场合抛出异常
        if (result.hasErrors()) {
            throw new IllegalArgumentException(
                    result.getError().getMessages().stream().map(WeaviateErrorMessage::getMessage).collect(joining("\n"))
            );
        }

        return resIds;
    }

    /**
     * modify by cgy: 增加元数据
     */
    private WeaviateObject buildObject(String id, Embedding embedding, String text, Metadata metadata) {
        Map<String, Object> props = new HashMap<>();

        if (metadata != null) {
            props.putAll(metadata.asMap());
        }

        props.put(METADATA_TEXT_SEGMENT, text);

        return WeaviateObject
                .builder()
                .className(objectClass)
                .id(id)
                .vector(embedding.vectorAsList().toArray(new Float[0]))
                .properties(props)
                .build();
    }


    private WeaviateObject buildObject(String id, Embedding embedding, String text) {
        Map<String, Object> props = new HashMap<>();
        props.put(METADATA_TEXT_SEGMENT, text);

        return WeaviateObject
                .builder()
                .className(objectClass)
                .id(id)
                .vector(embedding.vectorAsList().toArray(new Float[0]))
                .properties(props)
                .build();
    }

    private static EmbeddingMatch<TextSegment> toEmbeddingMatch(Map<String, ?> item) {
        Map<String, ?> additional = (Map<String, ?>) item.get(ADDITIONALS);
        String text = (String) item.get(METADATA_TEXT_SEGMENT);

        // modify by cgy: 返回数据增加元数据
        Map<String, String> meta = new HashMap<>();
        item.forEach((k, v) -> {
            if (v instanceof String && !METADATA_TEXT_SEGMENT.equals(k)) {
                meta.put(k, (String) v);
            }
        });
        TextSegment segment = null;
        if (text != null && !text.isEmpty()) {
            segment = TextSegment.from(text, Metadata.from(meta));
        }

        return new EmbeddingMatch<>(
                (Double) additional.get("certainty"),
                (String) additional.get("id"),
                Embedding.from(
                        ((List<Double>) additional.get("vector")).stream().map(Double::floatValue).collect(toList())
                ),
                segment
        );
    }
}
